"""
医学黑色素瘤检测应用

这个应用提供了一个简单的图形用户界面，用于上传图像并检测黑色素瘤区域。
使用Streamlit库构建，使用基于TensorFlow的U-Net模型进行分割。
"""

import streamlit as st
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt
from PIL import Image
import io
import tensorflow as tf
from utils import preprocess_image, predict_mask, overlay_mask
from unet_model import build_unet_model

# 设置页面标题和布局
st.set_page_config(page_title="医学黑色素瘤检测", layout="wide")

# 页面标题
st.title("医学黑色素瘤检测系统")
st.write("这个应用使用U-Net深度学习模型检测皮肤图像中的黑色素瘤区域。")

# 加载模型
@st.cache_resource
def load_model():
    """
    加载预训练模型（使用cache_resource避免重复加载）
    如果没有预训练模型，创建一个新模型（仅用于演示）
    """
    model_path = "models/melanoma_unet.h5"
    try:
        model = tf.keras.models.load_model(model_path)
        st.success("✅ 成功加载预训练模型")
    except:
        st.warning("⚠️ 未找到预训练模型，创建新模型（仅用于演示，需要训练才能获得准确结果）")
        model = build_unet_model(input_shape=(128, 128, 3))
        model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        
        # 创建models目录（如果不存在）
        os.makedirs("models", exist_ok=True)
        
    return model

# 添加侧边栏
st.sidebar.title("关于")
st.sidebar.info(
    """
    这个应用展示了如何使用深度学习模型（U-Net架构）对医学图像进行分割，
    以检测皮肤黑色素瘤区域。
    
    **使用方法**:
    1. 上传一张皮肤图像
    2. 点击"检测黑色素瘤"按钮
    3. 查看检测结果
    
    **技术栈**:
    - TensorFlow/Keras（深度学习框架）
    - U-Net（图像分割架构）
    - Streamlit（Web应用框架）
    - OpenCV（图像处理）
    """
)

# 创建主界面
col1, col2 = st.columns(2)

# 上传图像
with col1:
    st.header("上传图像")
    uploaded_file = st.file_uploader("选择一张皮肤图像...", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        # 显示上传的图像
        image = Image.open(uploaded_file)
        st.image(image, caption="上传的图像", use_container_width=True)
        
        # 转换为numpy数组
        image_np = np.array(image)
        
        # 检查图像是否为RGB
        if len(image_np.shape) == 2:  # 灰度图像
            image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
        elif image_np.shape[2] == 4:  # RGBA图像
            image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)

# 检测按钮和结果显示
with col2:
    st.header("检测结果")
    
    if uploaded_file is None:
        st.info("请先上传图像以进行黑色素瘤检测")
    else:
        detect_button = st.button("检测黑色素瘤")
        
        if detect_button:
            with st.spinner('正在检测...'):
                try:
                    # 加载模型
                    model = load_model()
                    
                    # 预处理图像
                    preprocessed_image = preprocess_image(image_np)
                    
                    # 预测掩码
                    pred_mask = model.predict(preprocessed_image)
                    pred_mask = pred_mask[0]  # 移除批次维度
                    
                    # 阈值化掩码
                    binary_mask = (pred_mask > 0.5).astype(np.uint8)
                    
                    # 将掩码叠加在原始图像上
                    overlaid_image = overlay_mask(image_np, binary_mask)
                    
                    # 显示结果
                    st.image(overlaid_image, caption="检测结果（红色区域表示黑色素瘤）", use_container_width=True)
                    
                    # 计算黑色素瘤覆盖率
                    melanoma_percentage = np.sum(binary_mask) / binary_mask.size * 100
                    st.info(f"黑色素瘤覆盖率: {melanoma_percentage:.2f}%")
                    
                    # 基于覆盖率的简单风险评估（仅作示范）
                    if melanoma_percentage > 10:
                        st.error("⚠️ 较高风险：检测到较大面积的黑色素瘤区域，建议咨询皮肤科医生")
                    elif melanoma_percentage > 2:
                        st.warning("⚗️ 中等风险：检测到黑色素瘤区域，建议进一步检查")
                    else:
                        st.success("✅ 低风险：检测到较小面积的黑色素瘤区域或未检测到")
                    
                    st.caption("注意：这只是一个演示应用，不应替代专业医疗诊断。")
                    
                except Exception as e:
                    st.error(f"检测过程中出现错误: {str(e)}")

# 添加模型训练部分（可选）
st.header("模型训练（演示）")
st.write("以下是模型训练过程的演示，真实应用中需要使用ISIC黑色素瘤数据集进行训练。")

col3, col4 = st.columns(2)

with col3:
    if st.button("准备训练数据"):
        with st.spinner("准备数据中..."):
            # 创建示例数据（仅用于演示）
            st.code("""
# 数据预处理示例代码
import numpy as np
from data_preprocessing import process_dataset

# 处理训练集
train_images, train_masks = process_dataset(
    images_dir="path/to/training/images",
    masks_dir="path/to/training/masks",
    output_dir="data",
    augment=True
)

# 示例输出
print(f"训练图像形状: {train_images.shape}")
print(f"训练掩码形状: {train_masks.shape}")
            """)
            st.success("数据准备完成！（这是一个演示，没有实际处理数据）")

with col4:
    if st.button("训练模型"):
        with st.spinner("训练模型中..."):
            # 显示训练过程（仅用于演示）
            st.code("""
# 模型训练示例代码
from train_model import train, plot_training_history

# 训练模型
model, history = train(
    train_images=train_images,
    train_masks=train_masks,
    test_images=test_images,
    test_masks=test_masks,
    batch_size=8,
    epochs=10,
    model_save_path="models/melanoma_unet.h5"
)

# 绘制训练历史
plot_training_history(history)
            """)
            
            # 显示模拟的训练进度
            progress_bar = st.progress(0)
            for i in range(101):
                # 更新进度条
                progress_bar.progress(i)
                if i % 10 == 0:
                    st.write(f"Epoch {i//10}/{10} - loss: {1.0-i/100:.4f} - accuracy: {0.5+i/200:.4f}")
                    
                if i == 100:
                    # 完成时显示模拟的训练曲线
                    fig, ax = plt.subplots(1, 2, figsize=(12, 4))
                    epochs = range(1, 11)
                    loss = [0.9, 0.7, 0.5, 0.4, 0.35, 0.32, 0.3, 0.28, 0.27, 0.26]
                    acc = [0.55, 0.65, 0.75, 0.8, 0.82, 0.84, 0.85, 0.86, 0.87, 0.88]
                    
                    ax[0].plot(epochs, loss)
                    ax[0].set_title('训练损失')
                    ax[0].set_xlabel('Epoch')
                    ax[0].set_ylabel('Loss')
                    
                    ax[1].plot(epochs, acc)
                    ax[1].set_title('训练准确率')
                    ax[1].set_xlabel('Epoch')
                    ax[1].set_ylabel('Accuracy')
                    
                    # 在Streamlit中显示图表
                    st.pyplot(fig)
                    
            st.success("模型训练完成！（这是一个演示，没有实际训练模型）")

# 添加项目信息和教育资源
st.header("黑色素瘤知识")
st.write("""
黑色素瘤是一种严重的皮肤癌，早期发现和诊断对治疗至关重要。以下是判断皮肤痣是否可能是黑色素瘤的"ABCDE法则"：

- **A（不对称性 Asymmetry）**：黑色素瘤通常形状不规则，若将痣分为两半，两边形状不一致。
- **B（边界 Border）**：黑色素瘤的边缘往往不规则或模糊。
- **C（颜色 Color）**：黑色素瘤常呈现多种颜色或颜色分布不均。
- **D（直径 Diameter）**：直径大于6mm的痣需要特别关注。
- **E（演变 Evolution）**：痣的大小、形状、颜色或高度的变化是重要警示信号。

深度学习模型可以帮助医生更准确地识别黑色素瘤，但应始终咨询专业医疗意见进行正式诊断。
""")

# 添加页脚
st.markdown("---")
st.caption("黑色素瘤检测系统 | 基于TensorFlow和U-Net架构 | 仅供教育和演示用途")
